#!/usr/bin/env python3
"""
验证ChromaDB安装状态
"""

import sys
import os

def check_chromadb_installation():
    """检查ChromaDB安装状态"""
    print("=== 检查ChromaDB安装状态 ===")
    
    try:
        # 尝试导入ChromaDB
        import chromadb
        print("✅ ChromaDB 已成功安装")
        print(f"   版本: {chromadb.__version__}")
        
        # 检查ChromaDB功能
        from chromadb import Client
        client = Client()
        print("✅ ChromaDB客户端创建成功")
        
        # 创建测试集合
        collection = client.create_collection("test_installation")
        print("✅ ChromaDB集合创建成功")
        
        # 添加测试文档
        collection.add(
            documents=["测试文档1", "测试文档2"],
            metadatas=[{"source": "test"}, {"source": "test"}],
            ids=["id1", "id2"]
        )
        print("✅ 测试文档添加成功")
        
        # 查询测试
        results = collection.query(query_texts=["测试"], n_results=2)
        print(f"✅ 查询测试成功，找到 {len(results['documents'][0])} 个结果")
        
        # 清理
        client.delete_collection("test_installation")
        print("✅ 测试清理完成")
        
        return True
        
    except ImportError:
        print("❌ ChromaDB 未安装")
        return False
    except Exception as e:
        print(f"❌ ChromaDB 功能测试失败: {e}")
        return False

def check_chromadb_integration():
    """检查ChromaDB与系统的集成"""
    print("\n=== 检查ChromaDB系统集成 ===")
    
    try:
        from utils.chroma_db import KnowledgeChromaDB, HAS_CHROMADB
        
        print(f"HAS_CHROMADB 状态: {HAS_CHROMADB}")
        
        if HAS_CHROMADB:
            print("✅ 系统检测到ChromaDB已安装")
            
            # 测试ChromaDB模式
            chroma_kb = KnowledgeChromaDB("integration_test")
            print(f"✅ ChromaDB知识库创建成功，模式: {chroma_kb.mode}")
            
            # 添加测试数据
            test_data = [
                {"question": "集成测试问题", "answer": "集成测试答案"}
            ]
            chroma_kb.batch_add_qa_pairs(test_data)
            print("✅ 测试数据添加成功")
            
            # 搜索测试
            results = chroma_kb.search_question("集成测试")
            print(f"✅ 搜索测试成功，找到 {len(results)} 个结果")
            
            # 统计信息
            stats = chroma_kb.get_stats()
            print(f"✅ 统计信息获取成功: {stats}")
            
            return True
        else:
            print("❌ 系统未检测到ChromaDB安装")
            return False
            
    except Exception as e:
        print(f"❌ ChromaDB集成测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def cleanup():
    """清理测试文件"""
    print("\n=== 清理测试文件 ===")
    
    # 清理ChromaDB测试数据
    try:
        import chromadb
        client = chromadb.Client()
        collections = client.list_collections()
        for collection in collections:
            if collection.name.startswith("test_"):
                client.delete_collection(collection.name)
                print(f"✅ 清理ChromaDB集合: {collection.name}")
    except:
        pass

if __name__ == "__main__":
    print("🔍 开始ChromaDB安装验证")
    
    # 运行测试
    installation_ok = check_chromadb_installation()
    integration_ok = check_chromadb_integration()
    
    # 清理
    cleanup()
    
    # 输出结果
    print("\n" + "="*60)
    print("📋 ChromaDB验证结果:")
    print(f"   安装状态: {'✅ 成功' if installation_ok else '❌ 失败'}")
    print(f"   系统集成: {'✅ 成功' if integration_ok else '❌ 失败'}")
    
    if installation_ok and integration_ok:
        print("\n🎉 ChromaDB已成功安装并集成到系统中！")
        print("💡 系统现在将使用ChromaDB作为向量数据库")
        sys.exit(0)
    else:
        print("\n❌ ChromaDB验证失败")
        sys.exit(1)